import os, sys, re, codecs
from pyexpat.errors import codes
import json
from collections import OrderedDict
import numpy as np
from torchaudio import kaldi_io
from kaldiio import WriteHelper


frame_shift = 0.01
frame_size = 0.25
sil_symbol = "!SIL"
batch_size = 16

class Jsonunit(object):
    """Storing word's start end 

    Args:
        object (_type_): _description_
    """
    def __init__(self, word, start, end):
        self.word = word
        self.start = start
        self.end = end
    
    def __str__(self) -> str:
        pass

    def get_endpoint(self) ->[float, float]:
        # comment: 
        return [float(self.start), float(self.end)]

    def get_word(self) ->str:
        # comment: 
        return self.word
    # end def


def timeseq2labels(sent_name, time_seq, total_time) -> (str, list, list, list):
    # comment: 
    if sent_name == "" or not isinstance(time_seq, list) or len(time_seq) == 0:
        return (sent_name, [], [], [])
    
    start_list, end_list = list(), list()
    label_list = []
    prev_end = -1
    for i in range(len(time_seq)):
        start, end = time_seq[i].get_endpoint()
        label = time_seq[i].get_word()
        if i == 0:
            if start > 0.0:
                start_list.append(0.0)
                end_list.append(start)
                label_list.append(0)
                start_list.append(start)
                end_list.append(end)
                if label != sil_symbol:
                    label_list.append(1)
                else:
                    label_list.append(0)
            else:
                start_list.append(start)
                end_list.append(end)
                if label != sil_symbol:
                    label_list.append(1)
                else:
                    label_list.append(0)
            prev_end = end
        elif i == len(time_seq) - 1:
            if end <= total_time:
                start_list.append(start)
                end_list.append(end)
                if label != sil_symbol:
                    label_list.append(1)
                else:
                    label_list.append(0)
            else:
                start_list.append(start)
                end_list.append(end)
                if label != sil_symbol:
                    label_list.append(1)
                else:
                    label_list.append(0)
            prev_end = end
        else:
            """use prev_end for time seq
            """
            if start > prev_end:
                # interval silence occur
                start_list.append(prev_end)
                end_list.append(start)
                label_list.append(0)
                start_list.append(start)
                end_list.append(end)
                if label != sil_symbol:
                    label_list.append(1)
                else:
                    label_list.append(0)
            else:
                start_list.append(start)
                end_list.append(end)
                if label != sil_symbol:
                    label_list.append(1)
                else:
                    label_list.append(0)
            prev_end = end
        
        print("{}\n{}\n{}\n".format(start_list, end_list, label_list))

    print('{} == {}, {} == {}\n'.format(len(start_list), len(end_list), len(end_list), len(label_list)))
    assert len(start_list) == len(end_list)
    assert len(end_list) == len(label_list)
    return (sent_name, start_list, end_list, label_list)
# end def


def read_utt2dur(arg4utt2dur):
    # comment: 
    if not os.path.exists(arg4utt2dur):
        raise FileNotFoundError
    
    dur_info = OrderedDict()
    with codecs.open(arg4utt2dur, "r") as f:
        sent_line = f.readline().strip()
        while sent_line:
            sent_id, dur = re.split(r"\s+", sent_line)
            dur_info[sent_id] = float(dur)
            sent_line = f.readline().strip()
    return dur_info
# end def


def calc_one_ratio_in_total(label_tensor):
    # comment: 
    if not isinstance(label_tensor, np.ndarray):
        raise TypeError("numpy array needed\n")

    if label_tensor.shape[0] == 0:
        return 0
    one_sum = np.sum(label_tensor, axis=0)
    return float(one_sum) / float(label_tensor.shape[0])
# end def


def main(ctm_file_name, utt2dur_fn, dst_dir):
    if (not os.path.exists(ctm_file_name)):
        # comment: 
        raise Exception("File not exits\n")
    
    dur_info = read_utt2dur(utt2dur_fn)
    time_info = OrderedDict()
    with open(ctm_file_name, "r") as f:
        prev_file_id = ''
        line = f.readline().strip("\n ")
        unk_found = False
        while line:
            line_units = re.split(r"\s+", line)
            if len(line_units) != 5:
                continue
            if line_units[4] == "<UNK>":
                line = f.readline().strip("\n")
                unk_found = True
                continue
            if prev_file_id == "":
                prev_file_id = line_units[0]
                time_info[prev_file_id] = []
                start = float(line_units[2])
                end = float(line_units[3]) + start
                time_info[prev_file_id].append(Jsonunit(line_units[4], start, end))
            else:
                if (prev_file_id != line_units[0]):
                    # comment: New sentence start
                    if unk_found:
                        del time_info[prev_file_id]
                        unk_found = False
                    prev_file_id = line_units[0]
                    time_info[prev_file_id] = []
                    start = float(line_units[2])
                    end = float(line_units[3]) + start
                    time_info[prev_file_id].append(Jsonunit(line_units[4], start, end))
                else:
                    # comment: continue this sentence
                    start = float(line_units[2])
                    end = float(line_units[3]) + start
                    time_info[prev_file_id].append(Jsonunit(line_units[4], start, end))
            
            line = f.readline().strip("\n ")

                # end if
    # end if
    # assert len(dur_info) == len(time_info)
    if not os.path.exists(dst_dir):
        os.mkdir(dst_dir)
    
    label_info = {}
    scp_or_ark_index = 0
    used_ids, used_numpy_array = [], []
    sent_id = 0
    for key, value in time_info.items():
        # comment: 
        sent_id += 1
        _, starts, ends, labels = timeseq2labels(key, value, total_time=dur_info[key])
        frame_index = 0
        label_info[key] = []
        for start, end, label in zip(starts, ends, labels):
            current_time = frame_index * frame_shift
            while current_time >= start and current_time < end and current_time < dur_info[key] - (frame_size-frame_shift):
                label_info[key].append(label)
                frame_index += 1
                current_time = frame_index * frame_shift
    
        format_label = np.array(label_info[key])
        one_ratio = calc_one_ratio_in_total(format_label)
        print("ratio:{}\nid:{}label:{}\n".format(one_ratio, key, format_label))
        if one_ratio > 1.1:
            continue
        
        used_ids.append(key)
        used_numpy_array.append(format_label)
        if len(used_ids) % batch_size == 0 or sent_id == len(time_info):
            scp_name = '{}/{}_{}.scp'.format(dst_dir, 'vad_label', scp_or_ark_index)
            ark_name = '{}/{}_{}.ark'.format(dst_dir, 'vad_label', scp_or_ark_index)
            kaldi_io_format = 'ark,scp:{},{}'.format(ark_name, scp_name)
            # start to write a sub scp and ark
            with WriteHelper(kaldi_io_format) as writer:
                for id, np_array in zip(used_ids, used_numpy_array):
                    # numpy matrix into this address
                    writer(id, np_array)
            del used_ids[:batch_size]
            del used_numpy_array[:batch_size]
            scp_or_ark_index += 1
        


    # end for
if __name__ == "__main__":
    if len(sys.argv) == 4:
        ctm_file_name = sys.argv[1]
        utt2dur = sys.argv[2]
        label_dir = sys.argv[3]
        main(ctm_file_name=ctm_file_name, utt2dur_fn=utt2dur, dst_dir=label_dir)
    else:
        sys.stdout.write("{} {} {}\n".format(sys.argv[0], "ctm", "label_dir"))
        raise Exception("Parameter is not enough")
